#!/usr/bin/env python3
#
# Copyright (c) 2022 Nordic Semiconductor ASA
#
# SPDX-License-Identifier: LicenseRef-Nordic-5-Clause

"""
Utility for manipulating DFU Multi Image packages.

DFU Multi Image package is a general-purpose update file consisting of
a CBOR-based header that describes contents of the package, followed by a number
of update components, such as firmware images for different MCU cores.

The CBOR header currently complies with the following format:
{
    "img": [
        {"id": 0, "size": 102400},
        {"id": 1, "size": 204800}
        ...
    ]
}

Usage examples:

Creating DFU Multi Image package:
./dfu_multi_image_tool.py create --image 0 app_update.bin --image 1 net_core_app_update.bin dfu_multi_image.bin

Showing DFU Multi Image package header:
./dfu_multi_image_tool.py show dfu_multi_image.bin
"""

import argparse
import os
import struct

import cbor2

# Buffer size used for file reads to ensure large files are not loaded into memory at once
READ_BUFFER_SIZE = 16 * 1024

def get_aligned_size_and_padding(non_aligned_size: int, align: int, last_image: bool) -> int:
    # Alignment on last image is not necessary
    if last_image or align <= 1:
        return non_aligned_size, 0

    residual = non_aligned_size % align
    needed_padding = (align - residual) % align
    return non_aligned_size + needed_padding, needed_padding


def generate_header(image: list, align: int) -> bytes:
    """
    Generate DFU Multi Image package header
    """

    image_data = []
    for id, path in image:
        image_size = os.path.getsize(path)
        aligned_size, _ = get_aligned_size_and_padding(image_size, align, id == image[-1][0])
        image_data.append({'id': int(id), 'size': aligned_size})

    header_data = {'img': image_data}
    header_cbor = cbor2.dumps(header_data)
    fixed_header_length = struct.calcsize('<H') # Will resolve to 2
    header_length_no_padding = fixed_header_length + len(header_cbor)

    _, padding_length = get_aligned_size_and_padding(header_length_no_padding, align, False)
    # As the format of the dfu_multi_image is fixed, the CBOR parser inside the dfu_multi_image
    # library will simply parse the part of the header that matches the expected format and
    # ignore any additional data.
    # The header length encoded in the 2 bytes before the CBOR data will account for the padding.
    # It can be done, as the CBOR parser treats this length as an upper limit, not the exact length.
    padding_data = bytes([0xff] * padding_length)

    return struct.pack('<H', len(header_cbor) + padding_length) +  header_cbor + padding_data


def parse_header(file: object) -> object:
    """
    Parse DFU Multi Image package header
    """

    header_fixed_size = struct.calcsize('<H')
    header_cbor_size, = struct.unpack('<H', file.read(header_fixed_size))
    header_cbor = file.read(header_cbor_size)

    return cbor2.loads(header_cbor)


def generate_image(images: list, align: int, output_file: str) -> None:
    """
    Generate DFU Multi Image package
    """

    with open(output_file, 'wb') as out_file:
        out_file.write(generate_header(images, align))

        for id, path in images:
            image_size = os.path.getsize(path)
            with open(path, 'rb') as file:
                while True:
                    chunk = file.read(READ_BUFFER_SIZE)
                    if not chunk:
                        break
                    out_file.write(chunk)
                _, padding_length = get_aligned_size_and_padding(image_size, align, id == images[-1][0])
                out_file.write(bytes([0xff] * padding_length))

def show_header(input_file: str) -> None:
    """
    Parse and print DFU Multi Image package header
    """

    with open(input_file, 'rb') as file:
        header = parse_header(file)
        print('Images:')

        for image in header['img']:
            print(f'- Id: {image["id"]}')
            print(f'  Size: {image["size"]}')


def main():
    parser = argparse.ArgumentParser(description='DFU Multi Image tool',
                                     fromfile_prefix_chars='@',
                                     allow_abbrev=False)
    subcommands = parser.add_subparsers(dest='subcommand', title='valid subcommands')

    create_parser = subcommands.add_parser(
        'create', help='Create DFU Multi Image package')
    create_parser.add_argument(
        '-i', '--image',
        required=True, action='append', nargs=2, metavar=('id', 'path'),
        help='Image to be included in package')
    create_parser.add_argument(
        '--align', type=int, default=1,
        help='Alignment of the start of every image. Gaps will be filled with 0xFF bytes.')
    create_parser.add_argument(
        'output_file', help='Path to output package file')

    show_parser = subcommands.add_parser(
        'show', help='Show DFU Multi Image package header')
    show_parser.add_argument(
        'input_file', help='Path to package file')

    args = parser.parse_args()

    if args.subcommand == 'create':
        generate_image(args.image, args.align, args.output_file)
    elif args.subcommand == 'show':
        show_header(args.input_file)
    else:
        parser.print_help()


if __name__ == "__main__":
    main()
